import torch
import paddle
from torch.utils.data import Dataset

# create Dataset
class create_Dataset(Dataset):
    def __init__(self,data, tokenizer, pinyin2id, tag2id):
        '''
        输入包括 
        1. text 句子 
        2. 拼音标注[[pos, pinyin],... ,[pos, pinyin]]
        3. 分词标注[B, I, O, E, S](O只出现在padding之后)
        4. 细粒度rhythm [0,1,0,...,1,0,1]
        5. 中粒度rhythm [0,1,0,...,1,0,1]
        6. 粗粒度rhythm [0,1,0,...,1,0,1]
        '''
        self.x = [tokenizer.convert_tokens_to_ids(i[0]) for i in data]
        self.pinyin = [[[j[0] for j in i[1][0]],[pinyin2id[j[1]] for j in i[1][0]]] for i in data]
        self.seg = [[tag2id[j] for j in i[1][1]] for i in data]
        self.rhythm1 = [i[1][2] for i in data]
        self.rhythm2 = [i[1][3] for i in data]
        self.rhythm3 = [i[1][4] for i in data]


    def __getitem__(self, index):
        return self.x[index],self.pinyin[index],self.seg[index],self.rhythm1[index],self.rhythm2[index],self.rhythm3[index]

    def __len__(self):
        return len(self.x)

class create_Dataset_type(Dataset):
    def __init__(self,data, tokenizer, pinyin2id, tag2id, type2id):
        '''
        输入包括 
        1. text 句子 
        2. 拼音标注[[pos, pinyin],... ,[pos, pinyin]]
        3. 分词标注[B, I, O, E, S](O只出现在padding之后)
        4. 细粒度rhythm [0,1,0,...,1,0,1]
        5. 中粒度rhythm [0,1,0,...,1,0,1]
        6. 粗粒度rhythm [0,1,0,...,1,0,1]
        '''
        self.x = [tokenizer.convert_tokens_to_ids(i[0]) for i in data]
        self.pinyin = [[[j[0] for j in i[1][0]],[pinyin2id[j[1]] for j in i[1][0]]] for i in data]
        self.seg = [[tag2id[j] for j in i[1][1]] for i in data]
        self.rhythm1 = [i[1][2] for i in data]
        self.rhythm2 = [i[1][3] for i in data]
        self.rhythm3 = [i[1][4] for i in data]
        self.type = [type2id[i[2]] for i in data]


    def __getitem__(self, index):
        return self.x[index],self.pinyin[index],self.seg[index],self.rhythm1[index],self.rhythm2[index],self.rhythm3[index],self.type[index]

    def __len__(self):
        return len(self.x)


class create_Dataset_single(Dataset):
    def __init__(self, data, tokenizer, data_type, target2id=None):
        '''
        方便测试时只输入部分label
        输入data_type包括 ['pinyin', 'seg', 'rhythm']
        1. text 句子 
        2. 拼音标注[[pos, pinyin],... ,[pos, pinyin]]
        3. 分词标注[B, I, O, E, S](O只出现在padding之后)
        4. 细粒度rhythm [0,1,0,...,1,0,1]
        5. 中粒度rhythm [0,1,0,...,1,0,1]
        6. 粗粒度rhythm [0,1,0,...,1,0,1]
        '''
        self.x = [tokenizer.convert_tokens_to_ids(i[0]) for i in data]
        if data_type=='polyphone':
            self.y = [[[j[0] for j in i[1]],[target2id[j[1]] for j in i[1]]] for i in data]
        elif data_type=='seg':
            self.y = [[target2id[j] for j in i[1]] for i in data]
        elif data_type=='rhythm':
            self.y1 = [i[1][0] for i in data]
            self.y2 = [i[1][1] for i in data]
            self.y3 = [i[1][2] for i in data]

        self.data_type = data_type

    def __getitem__(self, index):
        if self.data_type != 'rhythm':
            return self.x[index], self.y[index]
        else:
            return self.x[index], self.y1[index], self.y2[index], self.y3[index]

    def __len__(self):
        return len(self.x)


# dataloader ：padding 
def paddle_padding(data):
    x, y, z, r1, r2, r3 = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    z_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(z):
        end = x_lens[i]
        z_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    r1_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(r1):
        end = x_lens[i]
        r1_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    r2_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(r2):
        end = x_lens[i]
        r2_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    r3_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(r3):
        end = x_lens[i]
        r3_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')
        
    y_pad = None
    for i, y1 in enumerate(y):
        y1 = paddle.to_tensor(y1)
        y1 = paddle.concat([paddle.ones([1, y1.shape[1]], dtype='int64')*i, y1], axis=0)
        y_pad = y1 if y_pad is None else paddle.concat([y_pad, y1], axis=1)
    
    # y1 = paddle.to_tensor(y, dtype='int64')
    return x_pad, y_pad, z_pad, r1_pad, r2_pad, r3_pad


# dataloader ：padding 
def pytorch_padding(data):
    x, y, z, r1, r2, r3 = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = torch.LongTensor(s[:end])

    z_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(z):
        end = x_lens[i]
        z_pad[i, :end] = torch.LongTensor(s[:end])

    r1_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r1):
        end = x_lens[i]
        r1_pad[i, :end] = torch.LongTensor(s[:end])

    r2_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r2):
        end = x_lens[i]
        r2_pad[i, :end] = torch.LongTensor(s[:end])

    r3_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r3):
        end = x_lens[i]
        r3_pad[i, :end] = torch.LongTensor(s[:end])

    y_pad = None
    for i, y1 in enumerate(y):
        y1 = torch.LongTensor(y1)
        y1 = torch.concat([torch.ones([1, y1.shape[1]]).long()*i, y1], axis=0)
        y_pad = y1 if y_pad is None else torch.concat([y_pad, y1], axis=1)
    
    # y1 = paddle.to_tensor(y, dtype='int64')
    return x_pad, y_pad, z_pad, r1_pad, r2_pad, r3_pad

# dataloader ：padding with type 
def pytorch_padding_type(data):
    x, y, z, r1, r2, r3, t = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = torch.LongTensor(s[:end])

    z_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(z):
        end = x_lens[i]
        z_pad[i, :end] = torch.LongTensor(s[:end])

    r1_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r1):
        end = x_lens[i]
        r1_pad[i, :end] = torch.LongTensor(s[:end])

    r2_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r2):
        end = x_lens[i]
        r2_pad[i, :end] = torch.LongTensor(s[:end])

    r3_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r3):
        end = x_lens[i]
        r3_pad[i, :end] = torch.LongTensor(s[:end])

    y_pad = None
    for i, y1 in enumerate(y):
        y1 = torch.LongTensor(y1)
        y1 = torch.concat([torch.ones([1, y1.shape[1]]).long()*i, y1], axis=0)
        y_pad = y1 if y_pad is None else torch.concat([y_pad, y1], axis=1)
    
    return x_pad, y_pad, z_pad, r1_pad, r2_pad, r3_pad, torch.LongTensor(t)


def pytorch_seg_padding(data):
    x, z = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = torch.LongTensor(s[:end])
    
    z_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(z):
        end = x_lens[i]
        z_pad[i, :end] = torch.LongTensor(s[:end])

    return x_pad, z_pad


def paddle_seg_padding(data):
    x, z = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')
    
    z_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(z):
        end = x_lens[i]
        z_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    return x_pad, z_pad

    
def pytorch_rhythm_padding(data):
    x, r1, r2, r3 = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = torch.LongTensor(s[:end])
    
    r1_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r1):
        end = x_lens[i]
        r1_pad[i, :end] = torch.LongTensor(s[:end])

    r2_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r2):
        end = x_lens[i]
        r2_pad[i, :end] = torch.LongTensor(s[:end])

    r3_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(r3):
        end = x_lens[i]
        r3_pad[i, :end] = torch.LongTensor(s[:end])

    return x_pad, r1_pad, r2_pad, r3_pad

def paddle_rhythm_padding(data):
    x, r1, r2, r3 = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')
    
    r1_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(r1):
        end = x_lens[i]
        r1_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    r2_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(r2):
        end = x_lens[i]
        r2_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    r3_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(r3):
        end = x_lens[i]
        r3_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    return x_pad, r1_pad, r2_pad, r3_pad


def paddle_polyphone_padding(data):
    x, y = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = paddle.zeros([len(x), max(x_lens)], dtype='int64')
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = paddle.to_tensor(s[:end], dtype='int64')

    y_pad = None
    for i, y1 in enumerate(y):
        y1 = paddle.to_tensor(y1)
        y1 = paddle.concat([paddle.ones([1, y1.shape[1]], dtype='int64')*i, y1], axis=0)
        y_pad = y1 if y_pad is None else paddle.concat([y_pad, y1], axis=1)
    
    return x_pad, y_pad

def pytorch_polyphone_padding(data):
    x, y = zip(*data)
    
    x_lens = [len(t) for t in x]
    x_pad = torch.zeros([len(x), max(x_lens)]).long()
    for i, s in enumerate(x):
        end = x_lens[i]
        x_pad[i, :end] = torch.LongTensor(s[:end])

    
    y_pad = None
    for i, y1 in enumerate(y):
        y1 = torch.LongTensor(y1)
        y1 = torch.concat([torch.ones([1, y1.shape[1]]).long()*i, y1], axis=0)
        y_pad = y1 if y_pad is None else torch.concat([y_pad, y1], axis=1)
    
    return x_pad, y_pad